'''
FilePath: /jack/rockchip/facenet-pytorch/export_onnx.py
Copyright: 无锡宝通智能科技股份有限公司
Author: jiajunjie@boton-tech.com
LastEditTime: 2024-10-30 17:47:45
'''
import torch
import nets.facenet as models
model_path='./model_data/facenet_mobilenet.pth' #模型路径

model=models.Facenet(mode='predict')
device = torch.device('cpu')
model.load_state_dict(torch.load(model_path,map_location=device),strict=False) #模型加载
net=model.eval()
example=torch.rand(1,3,160,160) #给定一个输入
torch.onnx.export(model,(example),'./model_data/facenet_mobilenet.onnx',verbose=False,opset_version=12) #导出onnx模型
